Greptile Summary

This PR exposes two new runtime environment variables (NVTE_TENSOR_HANDLE_POOL_SIZE_MB and NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB) to make the internal handle pool sizes configurable, preserving the existing 20 MiB default and improving exhaustion error messages to include the configured size and the relevant env var name.

• Adds GetTensorHandlePoolSizeMB and GetTensorHandlePoolCapacity helpers in an anonymous namespace that read, validate, and convert env var values to pool capacities; the computed constants replace the previously hardcoded 20 * 1024 * 1024 / sizeof(T) expressions as member initializers in both allocator singletons.
• Updates docs/envvars.rst with documentation entries for the two new variables.

Confidence Score: 4/5

Safe to merge; the change is backward-compatible (default pool size unchanged) and the new code paths are straightforward.

The core logic is correct — validation guards against zero and overflow, the member-initializer ordering is right, and the singleton construction is safe. Two minor defects exist: partial string inputs like "50bad" are silently accepted as 50 due to no eof check in the shared getenv helper, and negative env var values wrap to SIZE_MAX and trigger the "too large" error rather than the intended "must be a positive integer" message. Both are input-validation edge cases with no correctness impact under documented usage.

transformer_engine/common/transformer_engine.cpp — specifically the env-var parsing and validation in GetTensorHandlePoolSizeMB

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant Env as Environment Variable
    participant Helper as GetTensorHandlePoolSizeMB
    participant Cap as GetTensorHandlePoolCapacity
    participant Alloc as TensorAllocator (singleton)
    participant App as Application Code

    App->>Alloc: first call to instance()
    activate Alloc
    Alloc->>Env: std::getenv("NVTE_TENSOR_HANDLE_POOL_SIZE_MB")
    Env-->>Helper: raw string value (or null default 20)
    Helper->>Helper: parse via getenv size_t
    Helper->>Helper: "NVTE_CHECK pool_size_mb > 0"
    Helper->>Helper: NVTE_CHECK pool_size_mb within bounds
    Helper-->>Cap: pool_size_mb
    Cap->>Cap: compute pool_size_bytes
    Cap->>Cap: "NVTE_CHECK pool_size_bytes >= sizeof(Tensor)"
    Cap-->>Alloc: MAX_TENSOR_NUM
    Alloc->>Alloc: memory.reserve(MAX_TENSOR_NUM)
    Alloc-->>App: singleton ready
    deactivate Alloc

    App->>Alloc: Allocate(mode, out, N)
    Alloc->>Alloc: "check available >= N"
    alt pool exhausted
        Alloc-->>App: NVTE_CHECK error with pool_size_mb and env var hint
    else space available
        Alloc-->>App: NVTETensor handles
    end

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

I am not opposed to creating such a variable, but I would really like to see an example of such legitimate use which goes over this limit. Could you run the experiment that is failing for you with https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/common/transformer_engine.cpp#L487 set to true and send me the log of that?